import tensorflow as tf
import os
from PIL import Image
import numpy as np
import pickle

os.environ["CUDA_VISIBLE_DEVICES"] = '1'


# The following functions can be used to convert a value to a type compatible
# with tf.Example.

def _bytes_feature(value):
    """Returns a bytes_list from a string / byte."""
    if isinstance(value, type(tf.constant(0))):
        value = value.numpy()  # BytesList won't unpack a string from an EagerTensor.
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


def _float_feature(value):
    """Returns a float_list from a float / double."""
    return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))


def _int64_feature(value):
    """Returns an int64_list from a bool / enum / int / uint."""
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))


def create_example(image_string, label):
    feature = {
        'label': _int64_feature(label),
        'image_raw': _bytes_feature(image_string),
    }

    return tf.train.Example(features=tf.train.Features(feature=feature))


def generate(imagelist, labels, record_file, IMG_SIZE):
    """使用PIllow 压缩，转numpy的bytes"""

    writer = tf.io.TFRecordWriter(record_file)
    num = 0
    for img, l in zip(imagelist, labels):
        try:
            with Image.open(img) as image_string:
                image_string = image_string.convert('RGB')
                image_string = image_string.resize((IMG_SIZE, IMG_SIZE), Image.ANTIALIAS)
                image_string = np.asarray(image_string, np.uint8)
                image_string = image_string.tobytes()
                tf_example = create_example(image_string, l)
                writer.write(tf_example.SerializeToString())
                num += 1
                if num % 1000 == 0:
                    print("has create {} example".format(num))
        except Exception as e:
            pass


def read_tf_record(serialized_example):
    image_feature_description = {
        'label': tf.io.FixedLenFeature([], tf.int64),
        'image_raw': tf.io.FixedLenFeature([], tf.string),
    }

    feature = tf.io.parse_example(serialized_example,
                                  features=image_feature_description)
    image = tf.io.decode_raw(feature['image_raw'], tf.uint8)
    image = tf.reshape(image, [224, 224, 3])
    image = tf.reverse(image, axis=[-1])  # RGB->BGR
    label = feature['label']

    return image, label


if __name__ == '__main__':
    root = '/hdd9/ppp/image_retrival'
    des_root = '/hdd9/ppp/image_retrival'

    with open(root + '/train.txt', 'r') as fp:
        images_anon = fp.readlines()

    images = []
    labels = []
    for img_anon in images_anon:
        img, label = img_anon.split('\t')
        images.append(img)
        labels.append(label.strip())
    classes = list(set(labels))
    Index = dict()
    for i, label in enumerate(classes):
        Index[label] = i
    labels = [Index[l] for l in labels]

    with open(des_root + '/cls.pk', 'wb') as fp:
        pickle.dump(Index, fp)

    record_file = des_root + '/train.record'
    generate(images, labels, record_file, IMG_SIZE=224)

    with open(root + '/test.txt', 'r') as fp:
        images_anon = fp.readlines()

    images = []
    labels = []
    for img_anon in images_anon:
        img, label = img_anon.split('\t')
        images.append(img)
        labels.append(label.strip())
    classes = list(set(labels))
    Index = dict()
    for i, label in enumerate(classes):
        Index[label] = i
    labels = [Index[l] for l in labels]
    record_file = des_root + '/test.record'
    generate(images, labels, record_file, IMG_SIZE=224)
